#!/bin/bash
# 默认变量
DEFAULT_CONDA_DIR="$HOME/.labelapp/miniconda"
DEFAULT_NODE_DIR="$HOME/.labelapp/nodejs"

PYTHON_VERSION="3.10"
CONDA_ENV_NAME="sam2"
REQUIREMENTS_FILE="requirements.txt"
MODEL_DIR="checkpoints"
PYTHON_CODE_DIR=$(pwd)
NODE_SCRIPT="../../configs/get_pytorch_command.js"

# 解析命令行参数
if [ "$1" == "local" ]; then
    echo "Using local Conda and Node.js installations."
    CONDA_CMD=conda
    ACTIVATE_CMD="conda run -n ${CONDA_ENV_NAME}"
    NODE_CMD=node
elif [ "$1" == "labelapp" ] || [ -z "$1" ]; then
    echo "Using labelapp Conda and Node.js installations."
    CONDA_DIR="$DEFAULT_CONDA_DIR"
    NODE_DIR="$DEFAULT_NODE_DIR"
    CONDA_CMD="${CONDA_DIR}/bin/conda"
    ACTIVATE_CMD="${CONDA_DIR}/bin/conda run -n ${CONDA_ENV_NAME}"
    NODE_CMD="${NODE_DIR}/bin/node"
else
    echo "Invalid argument. Usage: $0 [local|labelapp]"
    exit 1
fi

# 定义conda和nodejs默认的命令
echo "CONDA_CMD=${CONDA_CMD}"
echo "NODE_CMD=${NODE_CMD}"
echo "ACTIVATE_CMD=${ACTIVATE_CMD}"

# 检查 Conda 是否已安装
if ! command -v ${CONDA_CMD} &> /dev/null; then
    echo "Conda not found. Please install Conda first."
    exit 1
fi

# 获取 CUDA 版本
get_cuda_version() {
    # 尝试通过 nvcc 获取 CUDA 版本
    CUDA_VERSION=$(nvcc --version 2>/dev/null | grep "release" | awk '{print $6}' | cut -c2-)
    if [ -z "$CUDA_VERSION" ]; then
        # 如果 nvcc 不可用，尝试通过 nvidia-smi 获取 CUDA 版本
        CUDA_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n 1 | cut -d. -f1-2)
    fi
    # 只保留主版本和次版本（如 11.8.89 -> 11.8）
    echo "$CUDA_VERSION" | cut -d. -f1-2
}

# 检查 Conda 环境是否已存在
if ${CONDA_CMD} env list | grep -wq "^${CONDA_ENV_NAME}[[:space:]]"; then
    echo "Conda environment '$CONDA_ENV_NAME' already exists. Skipping creation."
else
    # 创建 Conda 环境
    echo "=== Creating Conda Environment ==="
    ${CONDA_CMD} create -n ${CONDA_ENV_NAME} python=${PYTHON_VERSION} -y
    if [ $? -ne 0 ]; then
        echo "Failed to create Conda environment. Exiting..."
        exit 1
    fi
fi

# 获取 CUDA 版本
CUDA_VERSION=$(get_cuda_version)
if [ -z "$CUDA_VERSION" ]; then
    echo "CUDA is not installed. Using CPU-only version..."
    CUDA_VERSION="cpu"
else
    echo "Detected CUDA version: ${CUDA_VERSION}"
fi

# 使用 Node.js 脚本获取 PyTorch 安装命令
if [ ! -f "$NODE_SCRIPT" ]; then
    echo "Node.js script $NODE_SCRIPT not found. Exiting..."
    exit 1
fi
INSTALL_COMMAND=$(${NODE_CMD} "$NODE_SCRIPT" "$CUDA_VERSION" "pip")
if [ $? -ne 0 ]; then
    echo "Failed to get PyTorch installation command. Exiting..."
    exit 1
fi

# 安装 PyTorch
echo "=== Installing PyTorch ==="
echo ${INSTALL_COMMAND}
${ACTIVATE_CMD} ${INSTALL_COMMAND}
if [ $? -ne 0 ]; then
    echo "Failed to install PyTorch. Exiting..."
    exit 1
fi

# 安装 requirements.txt 中的依赖
echo "=== Installing Dependencies ==="
${ACTIVATE_CMD} pip install -r ${PYTHON_CODE_DIR}/${REQUIREMENTS_FILE}
if [ $? -ne 0 ]; then
    echo "Failed to install dependencies. Exiting..."
    exit 1
fi

# 创建模型目录
echo "=== Creating Model Directory ==="
mkdir -p ${PYTHON_CODE_DIR}/${MODEL_DIR}
if [ $? -ne 0 ]; then
    echo "Failed to create model directory. Exiting..."
    exit 1
fi

# 下载模型文件
echo "=== Downloading Model Files ==="
MODEL_URLS=(
    "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt"
    "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt"
    "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt"
    "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt"
)

for URL in "${MODEL_URLS[@]}"; do
    FILENAME=$(basename ${URL})
    if [ -f "${PYTHON_CODE_DIR}/${MODEL_DIR}/${FILENAME}" ]; then
        echo "Model file ${FILENAME} already exists. Skipping download."
    else
        echo "Downloading ${FILENAME}..."
        wget -O ${PYTHON_CODE_DIR}/${MODEL_DIR}/${FILENAME} ${URL}
        if [ $? -ne 0 ]; then
            echo "Failed to download ${FILENAME}. Exiting..."
            exit 1
        fi
    fi
done

echo "=== Installation Completed Successfully ==="